"""
Memory Leak Detector Node for ComfyUI

Provides memory leak detection and analysis functionality.
"""

import time
import threading
from ..utils.memory_utils import MemoryTracker, get_memory_info, format_bytes


class MemoryLeakDetectorNode:
    """
    Memory Leak Detector Node - Detects and analyzes potential memory leaks
    """
    
    def __init__(self):
        self.tracker = MemoryTracker()
        self.monitoring_active = False
        self.monitoring_thread = None
        self.leak_history = []
        self.max_snapshots = 100  # Limit snapshots to prevent memory leaks
        self.max_leak_history = 50  # Limit leak history
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "action": (["Start Tracking", "Take Snapshot", "Detect Leaks", "Stop Tracking", "Get Report"], {
                    "default": "Take Snapshot"
                }),
                "snapshot_label": ("STRING", {
                    "default": "Snapshot",
                    "multiline": False
                }),
            },
            "optional": {
                "leak_threshold_mb": ("FLOAT", {
                    "default": 100.0,
                    "min": 10.0,
                    "max": 1000.0,
                    "step": 10.0,
                    "display": "slider"
                }),
                "auto_snapshot_interval": ("FLOAT", {
                    "default": 60.0,
                    "min": 10.0,
                    "max": 600.0,
                    "step": 10.0,
                    "display": "slider"
                }),
                "enable_auto_snapshots": ("BOOLEAN", {"default": False}),
            }
        }
    
    RETURN_TYPES = (
        "STRING",      # Action result
        "STRING",      # Leak report
        "INT",         # Number of leaks detected
        "BOOLEAN",     # Leaks found
        "STRING",      # Recommendations
    )
    
    RETURN_NAMES = (
        "action_result",
        "leak_report",
        "leak_count",
        "leaks_detected",
        "recommendations",
    )
    
    FUNCTION = "detect_memory_leaks"
    CATEGORY = "Memory Management"
    OUTPUT_NODE = True
    
    def detect_memory_leaks(self, action, snapshot_label="Snapshot", leak_threshold_mb=100.0, 
                          auto_snapshot_interval=60.0, enable_auto_snapshots=False):
        """Perform memory leak detection operations"""
        
        try:
            # Update tracker threshold
            self.tracker.leak_threshold = leak_threshold_mb * 1024 * 1024  # Convert MB to bytes
            
            if action == "Start Tracking":
                return self._start_tracking(enable_auto_snapshots, auto_snapshot_interval)
            
            elif action == "Take Snapshot":
                return self._take_snapshot(snapshot_label)
            
            elif action == "Detect Leaks":
                return self._detect_leaks()
            
            elif action == "Stop Tracking":
                return self._stop_tracking()
            
            elif action == "Get Report":
                return self._get_comprehensive_report()
            
            else:
                return (
                    f"Unknown action: {action}",
                    "No leak report available",
                    0,
                    False,
                    "Select a valid action"
                )
                
        except Exception as e:
            error_msg = f"Memory leak detection error: {str(e)}"
            return (
                error_msg,
                error_msg,
                0,
                False,
                "Fix the error and try again"
            )
    
    def _start_tracking(self, enable_auto_snapshots, auto_snapshot_interval):
        """Start memory leak tracking"""
        
        try:
            self.tracker.start_tracking()
            
            # Take initial snapshot
            self.tracker.take_snapshot("Initial_Baseline")
            
            # Start auto-snapshots if enabled
            if enable_auto_snapshots:
                self._start_auto_snapshots(auto_snapshot_interval)
            
            result = "✅ Memory leak tracking started successfully"
            report = self._create_tracking_status_report(True, enable_auto_snapshots, auto_snapshot_interval)
            
            return (
                result,
                report,
                0,
                False,
                "Take snapshots during your workflow to detect leaks"
            )
            
        except Exception as e:
            result = f"❌ Failed to start tracking: {str(e)}"
            return (
                result,
                result,
                0,
                False,
                "Check system permissions and memory availability"
            )
    
    def _take_snapshot(self, label):
        """Take a memory snapshot"""
        
        if not hasattr(self.tracker, 'snapshots') or not self.tracker.snapshots:
            return (
                "❌ Tracking not started - use 'Start Tracking' first",
                "No tracking session active",
                0,
                False,
                "Start tracking before taking snapshots"
            )
        
        snapshot = self.tracker.take_snapshot(label)
        
        # Enforce snapshot limits to prevent memory leaks
        if len(self.tracker.snapshots) > self.max_snapshots:
            # Remove oldest snapshots, keeping baseline
            snapshots_to_remove = len(self.tracker.snapshots) - self.max_snapshots
            self.tracker.snapshots = self.tracker.snapshots[:1] + self.tracker.snapshots[1+snapshots_to_remove:]
        
        if snapshot:
            memory_info = get_memory_info()
            current_usage = memory_info.get('process_rss', 0)
            
            result = f"✅ Snapshot '{label}' taken successfully"
            report = self._create_snapshot_report(label, current_usage, len(self.tracker.snapshots))
            
            return (
                result,
                report,
                0,
                False,
                "Continue workflow and take more snapshots to detect leaks"
            )
        else:
            result = "❌ Failed to take snapshot - tracking may not be active"
            return (
                result,
                result,
                0,
                False,
                "Ensure tracking is started before taking snapshots"
            )
    
    def _detect_leaks(self):
        """Detect memory leaks from snapshots"""
        
        if len(self.tracker.snapshots) < 2:
            return (
                "❌ Need at least 2 snapshots to detect leaks",
                "Insufficient snapshots for leak detection",
                0,
                False,
                "Take more snapshots during your workflow"
            )
        
        leaks = self.tracker.detect_leaks()
        leak_count = len(leaks)
        leaks_detected = leak_count > 0
        
        # Add to leak history
        self.leak_history.append({
            'timestamp': time.time(),
            'leak_count': leak_count,
            'leaks': leaks
        })
        
        result = f"🔍 Leak detection completed: {leak_count} potential leaks found"
        report = self._create_leak_report(leaks)
        recommendations = self._generate_leak_recommendations(leaks)
        
        return (
            result,
            report,
            leak_count,
            leaks_detected,
            recommendations
        )
    
    def _stop_tracking(self):
        """Stop memory leak tracking"""
        
        try:
            self.tracker.stop_tracking()
            self._stop_auto_snapshots()
            
            total_snapshots = len(self.tracker.snapshots)
            total_leaks = sum(len(entry['leaks']) for entry in self.leak_history)
            
            result = f"🛑 Memory leak tracking stopped"
            report = self._create_stop_report(total_snapshots, total_leaks)
            
            return (
                result,
                report,
                total_leaks,
                total_leaks > 0,
                "Review the leak history and address any identified issues"
            )
            
        except Exception as e:
            result = f"❌ Error stopping tracking: {str(e)}"
            return (
                result,
                result,
                0,
                False,
                "Tracking may have already been stopped"
            )
    
    def _get_comprehensive_report(self):
        """Get comprehensive memory leak report"""
        
        report = self._create_comprehensive_report()
        total_leaks = sum(len(entry['leaks']) for entry in self.leak_history)
        
        return (
            "📊 Comprehensive leak report generated",
            report,
            total_leaks,
            total_leaks > 0,
            "Review the complete analysis and take action on identified leaks"
        )
    
    def _create_tracking_status_report(self, tracking_active, auto_snapshots, interval):
        """Create tracking status report"""
        
        report_lines = ["=== MEMORY LEAK TRACKING STATUS ===\n"]
        
        report_lines.append(f"Tracking Status: {'ACTIVE' if tracking_active else 'INACTIVE'}")
        report_lines.append(f"Auto Snapshots: {'ENABLED' if auto_snapshots else 'DISABLED'}")
        if auto_snapshots:
            report_lines.append(f"Snapshot Interval: {interval:.1f} seconds")
        
        report_lines.append(f"Leak Threshold: {format_bytes(self.tracker.leak_threshold)}")
        report_lines.append(f"Snapshots Taken: {len(self.tracker.snapshots)}")
        
        return "\n".join(report_lines)
    
    def _create_snapshot_report(self, label, memory_usage, snapshot_count):
        """Create snapshot report"""
        
        report_lines = ["=== MEMORY SNAPSHOT REPORT ===\n"]
        
        report_lines.append(f"Snapshot Label: {label}")
        report_lines.append(f"Current Memory Usage: {format_bytes(memory_usage)}")
        report_lines.append(f"Total Snapshots: {snapshot_count}")
        report_lines.append(f"Timestamp: {time.strftime('%Y-%m-%d %H:%M:%S')}")
        
        return "\n".join(report_lines)
    
    def _create_leak_report(self, leaks):
        """Create detailed leak report"""
        
        if not leaks:
            return "✅ No memory leaks detected in the current analysis"
        
        report_lines = ["=== MEMORY LEAK DETECTION REPORT ===\n"]
        
        report_lines.append(f"🚨 {len(leaks)} potential memory leaks detected:\n")
        
        for i, leak in enumerate(leaks, 1):
            report_lines.append(f"LEAK #{i}:")
            report_lines.append(f"  File: {leak.get('filename', 'Unknown')}")
            report_lines.append(f"  Size Increase: {format_bytes(leak.get('size_diff', 0))}")
            report_lines.append(f"  Object Count Change: {leak.get('count_diff', 0)}")
            report_lines.append(f"  Size (MB): {leak.get('size_diff_mb', 0):.2f}")
            report_lines.append("")
        
        return "\n".join(report_lines)
    
    def _create_stop_report(self, total_snapshots, total_leaks):
        """Create tracking stop report"""
        
        report_lines = ["=== TRACKING SESSION SUMMARY ===\n"]
        
        report_lines.append(f"Total Snapshots Taken: {total_snapshots}")
        report_lines.append(f"Total Leaks Detected: {total_leaks}")
        report_lines.append(f"Leak Detection Sessions: {len(self.leak_history)}")
        
        if self.leak_history:
            report_lines.append(f"\nLeak History:")
            for i, entry in enumerate(self.leak_history, 1):
                timestamp = time.strftime('%H:%M:%S', time.localtime(entry['timestamp']))
                report_lines.append(f"  Session {i} ({timestamp}): {entry['leak_count']} leaks")
        
        return "\n".join(report_lines)
    
    def _create_comprehensive_report(self):
        """Create comprehensive leak analysis report"""
        
        report_lines = ["=== COMPREHENSIVE MEMORY LEAK ANALYSIS ===\n"]
        
        # Tracking status
        is_tracking = hasattr(self.tracker, 'snapshots') and len(self.tracker.snapshots) > 0
        report_lines.append(f"Tracking Status: {'ACTIVE' if is_tracking else 'INACTIVE'}")
        
        if is_tracking:
            report_lines.append(f"Total Snapshots: {len(self.tracker.snapshots)}")
            report_lines.append(f"Leak Threshold: {format_bytes(self.tracker.leak_threshold)}")
        
        # Current memory state
        memory_info = get_memory_info()
        report_lines.append(f"\nCURRENT MEMORY STATE:")
        report_lines.append(f"  Process RSS: {format_bytes(memory_info.get('process_rss', 0))}")
        report_lines.append(f"  System Usage: {memory_info.get('system_percent', 0):.1f}%")
        
        # Leak history analysis
        if self.leak_history:
            total_leaks = sum(entry['leak_count'] for entry in self.leak_history)
            avg_leaks = total_leaks / len(self.leak_history)
            
            report_lines.append(f"\nLEAK ANALYSIS:")
            report_lines.append(f"  Detection Sessions: {len(self.leak_history)}")
            report_lines.append(f"  Total Leaks Found: {total_leaks}")
            report_lines.append(f"  Average Leaks/Session: {avg_leaks:.1f}")
            
            # Most recent leaks
            if self.leak_history[-1]['leaks']:
                report_lines.append(f"\nMOST RECENT LEAKS:")
                for leak in self.leak_history[-1]['leaks'][:3]:  # Top 3
                    report_lines.append(f"  • {leak.get('filename', 'Unknown')}: {format_bytes(leak.get('size_diff', 0))}")
        else:
            report_lines.append(f"\nNo leak detection history available")
        
        return "\n".join(report_lines)
    
    def _generate_leak_recommendations(self, leaks):
        """Generate recommendations based on detected leaks"""
        
        if not leaks:
            return "No leaks detected - memory usage appears stable"
        
        recommendations = []
        
        # Analyze leak patterns
        large_leaks = [leak for leak in leaks if leak.get('size_diff_mb', 0) > 500]  # > 500MB
        file_leaks = {}
        
        for leak in leaks:
            filename = leak.get('filename', 'Unknown')
            if filename not in file_leaks:
                file_leaks[filename] = 0
            file_leaks[filename] += leak.get('size_diff_mb', 0)
        
        # Generate specific recommendations
        if large_leaks:
            recommendations.append(f"🚨 {len(large_leaks)} large memory leaks detected - immediate action required")
        
        if len(leaks) > 5:
            recommendations.append("Multiple leaks detected - review memory management practices")
        
        # File-specific recommendations
        top_leak_file = max(file_leaks.items(), key=lambda x: x[1]) if file_leaks else None
        if top_leak_file and top_leak_file[1] > 100:  # > 100MB
            recommendations.append(f"Primary leak source: {top_leak_file[0]} ({top_leak_file[1]:.1f}MB)")
        
        recommendations.append("Consider using aggressive memory cleanup")
        recommendations.append("Monitor workflow for memory-intensive operations")
        
        return " | ".join(recommendations)
    
    def _start_auto_snapshots(self, interval):
        """Start automatic snapshot thread"""
        
        if self.monitoring_active:
            return
        
        self.monitoring_active = True
        self.monitoring_thread = threading.Thread(
            target=self._auto_snapshot_loop, 
            args=(interval,), 
            daemon=True
        )
        self.monitoring_thread.start()
    
    def _stop_auto_snapshots(self):
        """Stop automatic snapshot thread"""
        
        self.monitoring_active = False
        if self.monitoring_thread:
            self.monitoring_thread.join(timeout=5.0)
    
    def _auto_snapshot_loop(self, interval):
        """Automatic snapshot loop"""
        
        snapshot_count = 0
        
        while self.monitoring_active:
            try:
                time.sleep(interval)
                if self.monitoring_active:
                    snapshot_count += 1
                    label = f"Auto_{snapshot_count}"
                    self.tracker.take_snapshot(label)
            except Exception as e:
                print(f"Auto-snapshot error: {e}")
                time.sleep(interval) 